package model

import (
	"context"
	"sync"
)

type TaskMap struct {
	lock     sync.RWMutex
	taskWork map[string]*TaskWorker
}

func NewTaskMap() *TaskMap {
	return &TaskMap{
		taskWork: make(map[string]*TaskWorker),
	}
}

func (t *TaskMap) GetTaskWork(taskId string) *TaskWorker {
	t.lock.RLock()
	defer t.lock.RUnlock()
	return t.taskWork[taskId]
}

func (t *TaskMap) SetTaskWork(taskId string, worker *TaskWorker) {
	t.lock.Lock()
	defer t.lock.Unlock()
	t.taskWork[taskId] = worker
}

func (t *TaskMap) DelTaskWork(taskId string) {
	t.lock.Lock()
	defer t.lock.Unlock()
	delete(t.taskWork, taskId)
}

type TaskWorker struct {
	TaskId  string
	Context context.Context
	Cancel  context.CancelFunc
}
